import onnx
from onnx import helper
from onnx import TensorProto
import numpy as np

# reshape
inp = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1,6,3,6])
shape = helper.make_tensor_value_info('shape', TensorProto.INT64, [5])
shape_weight = np.array([1, 3, 2, 3, 6])
const_shape = helper.make_tensor('shape', TensorProto.INT64, [5], shape_weight)

outp = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1, 3, 2, 3, 6])
reshape_def = helper.make_node(
    'Reshape',
    ['X', 'shape'],
    ['Y']
)

# transpose
transpose_output = helper.make_tensor_value_info('transpose_output', TensorProto.FLOAT, [1, 2, 3, 3, 6])
transpose_def = helper.make_node(
    'Transpose',
    ['Y'],
    ['transpose_output'],
    perm=[0, 2, 1, 3, 4]
)

# reshape 2
shape2 = helper.make_tensor_value_info('shape2', TensorProto.INT64, [4])
shape_weight2 = np.array([1, 6, 3, 6])
const_shape2 = helper.make_tensor('shape2', TensorProto.INT64, [4], shape_weight2)

reshape2_output = helper.make_tensor_value_info('reshape2_output', TensorProto.FLOAT, [1, 6, 3, 6])

reshape_def2 = helper.make_node(
    'Reshape',
    ['transpose_output', 'shape2'],
    ['reshape2_output'],
)

graph_def = helper.make_graph(
    [reshape_def, transpose_def, reshape_def2],
    "test_reshape_model",
    [inp, shape],
    [reshape2_output],
    initializer=[const_shape, const_shape2]
)




mode_def = helper.make_model(graph_def, producer_name='onnx-example')
onnx.checker.check_model(mode_def)
onnx.save(mode_def, "./shuffle.onnx")